import collections
import functools
import operator
import re
import tokenize


def query(obj, path, ret_keys=True, str_keys=True):
    """
    Return the value in a given JSON path of 'obj'. 
    
    Arguments
    ---------
    
    obj : JSON-like object 
        JSON structure (list-like or dictionary-like)   
    path : str, iterable
        The path can be a string of the form "$.child.toys.0" or an iterable 
        that expands to ['child', 'toys', 0]
    dict_query : bool
        Return this value if obj does not have the given path. 
        
    Examples
    --------
    
    Given the JSON-like structure
    
    >>> obj = { "store": {
    ...     "book": [ 
    ...       { "category": "reference",
    ...         "author": "Nigel Rees",
    ...         "title": "Sayings of the Century",
    ...         "discount price": 8.95
    ...       },
    ...       { "category": "fiction",
    ...         "author.name": "Evelyn",
    ...         "title": "Sword of Honour",
    ...         "price": 12.99
    ...       }
    ...     ]
    ...   }
    ... }
    
    The path can be an iterable such as a Python list

    >>> has_price = lambda L: [ (i, x) for (i, x) in enumerate(L) if 'price' in x ]
    >>> query(obj, ['store', 'book', has_price, 0, 'title'])
    { '$.store.book.1.title': 'Sword of Honour' }

    Or may also use a convenient string syntax: 
    
    >>> query(obj, '&.store.book[?(price)].title')
    { '$.store.book.1.title': 'Sword of Honour' }
    
     
    We can match one or an arbitrary number of levels using one or two 
    asterisks.
      
    >>> query(obj, '&.store.book.*.title')
    { '$.store.book.0.title': 'Sayings of the Century', '$.store.book.1.title': 'Sword of Honour' }
    
    >>> query(obj, '&.**.title')
    { '$.store.book.0.title': 'Sayings of the Century', '$.store.book.1.title': 'Sword of Honour' }
    
    
    Missing objects return an empty dict.
    
    >>> query(obj, '&.foo.bar', None)
    None
    >>> getitem(obj, '&.foo.bar')
    Traceback
    ...
    asdasd
    
    
    The leading $ is not optional:
    
    >>> getitem(obj, 'store.book')
    Traceback
    ...
    ValueError
    
    
    Observations
    ------------
    
    This standard is vaguely based on the notation defined at
    http://goessner.net/articles/JsonPath/.
    """

    resp_dict = {} #TODO: ordered dict
    curr_obj = obj
    for k in _iter_query_string(path, True):
        curr_obj = curr_obj[k]
    return curr_obj

def _iter_query_string(str, has_at):
    return iter(str)



class ArrayFilter(object):
    def __call__(self, array):
        raise NotImplementedError

class FilterHasItem(ArrayFilter):
    def __init__(self, item):
        self.item = item

    def __call__(self, obj):
        item = self.item
        return ((k, v) for (k, v) in str_itemize(obj) if hasitem(v, item))

class FilterBinOp(ArrayFilter):
    def __init__(self, item, op, value, value_right=True):
        self.item = item
        if value_right:
            self.op = lambda x: op(x, value)
        else:
            self.op = functools.partial(op, value)

    def __call__(self, obj):
        item = self.item
        op = self.op
        return ((k, v) for (k, v) in str_itemize(obj) if op(getitem(v, item)))

class Slice(ArrayFilter):
    pass

class CSV(ArrayFilter):
    pass

class Foo:
    def array_operation(self, obj, array_op):
        array_op = array_op.replace(' ', '')

        # filter
        if array_op.startswith('?(') and array_op.endswith(')'):
            filter = array_op[2:-1]
            return self.array_filter(obj, filter)

        # sort
        if array_op.startswith('#(') and array_op.endswith(')'):
            key = array_op[2:-1]
            return self.array_sort(obj, key)

        # select by list of indexes
        m = re_cvs.match(array_op)
        if m:
            idxs = map(int, array_op.split(','))
            return [ obj[idx] for idx in idxs ]

        # slice operations
        if ':' in array_op:
            slice = [ (int(x) if x else None) for x in array_op.split(':') ]
            if len(slice) == 2:
                a, b = slice
                return obj[a:b]
            elif len(slice) == 3:
                a, b, c = slice
                return obj[a:b:c]
            else:
                raise IndexError("Invalid slice '%s' at %s" % (array_op, self.orig_path))

        raise IndexError("unsupported array index '%s' in path '%s'" % (array_op, self.orig_path))

    def array_filter(self, obj, filter_):
        """filter array given by obj"""
        try:
            key, op, num, st = re_filter.findall(filter_)[0]
        except IndexError:
            raise KeyError("invalid array filter: ?(%s)" % filter_)
        keys = key.split('.')

        # define the getter function
        key = key.lstrip('@')
        walk = self.walk
        def getitem(x):
            try:
                obj, p = walk(x, key)
                assert not p, 'invalid traling path: %s' % p
                return obj
            except KeyError:
                return None

        # simple filter: check if properties exist
        if not op:
            f = lambda x: not (getitem(x) is None)
            return filter(f, obj)

        # define acceptance criteria
        op = filter_operators[op]
        if num:
            try:
                other = int(num)
            except ValueError:
                other = float(num)
        elif st:
            other = st[1:-1].decode('string_escape')
        else:
            raise KeyError("invalid filter '%s'" % filter)

        # check which elements in the array satisfy the acceptance criteria
        def accept(x):
            item = getitem(x)
            if not other is None:
                return op(item, other)
        return filter(accept, obj)

    def array_sort(self, obj, key):
        """filter array given by obj"""
        keys = key.split('.')

        # define the getter function
        key = key.lstrip('@')
        walk = self.walk
        def getitem(x):
            try:
                obj, p = walk(x, key)
                assert not p, 'invalid traling path: %s' % p
                return obj
            except KeyError:
                return None

        # simple filter: check if properties exist
        items = [ (x, getitem(x)) for x in obj ]
        items.sort(key=lambda x: x[1])
        nulls = [ x[0] for x in items if x[1] is None ]
        non_nulls = [ x[0] for x in items if not x[1] is None ]
        return non_nulls + nulls

    def recursive_descent(self, obj, path):
        visited = self.visited
        qset = self.queryset
        if id(obj) in visited:
            return
        else:
            visited.append(id(obj))

        if isinstance(obj, (basestring, int, float, bool, type(None))):
            return

        if isinstance(obj, collections.Mapping):
            try:
                self._rdescent_item(obj, path)
            except:
                pass
            obj = obj.values()

        if isinstance(obj, collections.Sequence):
            for item in obj:
                try:
                    self._rdescent_item(obj, path)
                except:
                    pass
                self.recursive_descent(item, path)

    def _rdescent_item(self, obj, path):
        item = self(obj, path)
        if not id(item) in self.visited:
            self.visited.append(id(item))
            if isinstance(item, queryset):
                self.queryset.extend(item)
            else:
                self.queryset.append(item)

    def expand_asterisk(self, obj, path):
        if isinstance(obj, collections.Mapping):
            seq = obj.items()
        else:
            seq = iter(obj)

        expansions = (self(o, path) for o in obj)
        result = queryset()
        for item in obj:
            try:
                item = self(item, path)
            except (KeyError, IndexError):
                continue
            if isinstance(item, queryset):
                result.extend(item)
            else:
                result.append(item)
        return result

    def __call__(self, obj, path):
        path = path.strip()
        if not path:
            return obj
        elif path.startswith('.*'):
            path = path[2:]
            return self.expand_asterisk(obj, path)
        elif path.startswith('..'):
            path = path[1:]
            self.visited = set()
            self.queryset = []
            self.recursive_descent(obj, path)
            return self.queryset
        else:
            # simple walk to object
            obj, path = self.walk(obj, path)
            if path:
                return self(obj, path)
            else:
                return obj

def jpath_update(json, updates):
    for p, v in updates.items():
        jpath_set(json, p, v)

#===============================================================================
#                   Helper functions: REGEX and constants
#===============================================================================
class queryset(list): pass

re_cvs = re.compile('^\d+\s*(,\d+\s*)*\s*\Z')
re_sort = re.compile(r'^\s*#.\s*([a-zA-Z]\w*)')
re_filter = re.compile(r'^\s*(@.(?:[a-zA-Z]\w*))\s*(==|!=|>=|<=|<|>|is)?\s*(%s)?%s?\s*\Z' %
                    (tokenize.Number.replace('[lL]+', '').replace('(', '(?:'),
                     tokenize.String))

filter_operators = {
    '==': operator.eq,
    '!=': operator.ne,
    '>=': operator.ge,
    '<=': operator.le,
    '>': operator.gt,
    '<': operator.lt, }

#===============================================================================
#                                TEST
#===============================================================================
if __name__ == '__main__':
    import doctest
    doctest.testmod(optionflags=doctest.REPORT_ONLY_FIRST_FAILURE, verbose=0)
